from math import inf
import numpy as np
from sympy import *
from scipy import optimize


def getData():
    _x1, _x2 = symbols('x1 x2')
    _f_x = _x1**2 + _x2**2 + 6*_x1 + 9
    _x = np.array((2, 0)).reshape((2, 1))
    _A = np.array([[2,1],
                   [1,0],
                   [0,1]])
    _b = np.array([4, 0, 0]).reshape((3, 1))
    _E=np.empty((0,0))
    return {
        'f_x': _f_x,
        'A': _A,
        'E':_E,
        'x': _x,
        'b': _b,
    }


def list2dict(x: list):
    _dict = dict()
    for i in range(len(x)):
       _dict['x{}'.format(i + 1)] = x[i][0]
    return _dict

def Simplex(f_x, A, b, s):
    '''
    单纯形法解线性方程
    :param f_x: 线性函数的系数矩阵
    :param A: 约束不等式的系数矩阵
    :param b: 约束不等式的右值矩阵
    :param s: 参数的范围矩阵
    :return: 最优解
    :testData:
        f_x = np.array([4, -1])
        A = np.array([[-1, 1], [-1, -1]])
        b = np.array([5, 0])
        s= ((None, 3),(None, None))
        print(Simplex(f_x, A, b, s))
    '''
    return optimize.linprog(f_x, A, b, bounds=s)


def correctX(x, acc=7):
    for i in range(len(x)):
        x[i][0] = round(x[i][0], acc)
    return x


def getMinPoint(f_lambda, MaxLambda, x):
    '''
    :param f_lambda:目标函数被lambda更新后变为的关于lambda的函数
    :param MaxLambda:函数的定义域为[0,MaxLambda]
    :param x:函数自变量的标识
    :return:函数达到最小值时所达到的点
    该方法的思想，目标函数经lambda转化后，变成只有一个自变量的函数，将函数关于该自变量求导，得到驻点，在驻点和端点的集合中找到函数值最小时的自变量取值。
    '''
    df_lambda = diff(f_lambda)
    _candidate = solve(df_lambda, x)
    _candidate.append(MaxLambda)
    _candidate = np.array(_candidate)
    _candidate = _candidate[_candidate <= MaxLambda]
    _res = 0
    _subs = {'l': 0}
    _function = f_lambda.evalf(subs=_subs)
    for _t in _candidate:
        _subs['l'] = _t
        _temp = f_lambda.evalf(subs=_subs)
        if _temp < _function:
            _function = _temp
            _res = _t
    return _res


class ZouUtils:

    def saperateMatrix(A, x, b):
        '''
        :param A:不等式约束的系数矩阵A.shape=m*n
        :param x:当前状态下的自变量取值x.shape=n*1
        :param b:不等式约束的右值b.shape=m*1
        :return:字典，拆分后的A,b,以及下标
        :testData:
            x=np.array([1,2]).reshape((2,1))
            A=np.array([[1,5],
                    [1,2],
                    [2,1],
                    [0,3],
                    [-1,5]])
            b=np.array([11,5,2,2,9]).reshape((5,1))
            print(np.matmul(A,x))
            print(ZouUtils.saperateMatrix(A, x, b))
        '''
        _b = np.matmul(A, x)
        _index = np.equal(_b, b).T
        e_index = np.array(np.where(_index == True))[1, :]
        g_index = np.array(np.where(_index == False))[1, :]
        _A1 = A[e_index]
        _b1 = b[e_index]
        _A2 = A[g_index]
        _b2 = b[g_index]

        return {'A1': _A1,
                'A2': _A2,
                'b1': _b1,
                'b2': _b2,
                'EI': e_index,
                'GI': g_index}

    def getGradient(f_x, x):
        '''
        :param f_x:目标函数的sympy表达式
        :param x: 自变量取值x.shape=n*1
        :return:梯度
        :testData:
            x=np.array([1,2]).reshape((2,1))
            x1, x2 = symbols('x1 x2')
            Y = x1 ** 2 - x2 ** 2
            print(ZouUtils.getGradient(Y,x))
        '''
        nums = max(x.shape)
        _xs = list()
        _subdic = dict()
        _delta_f = list()
        for i in range(nums):
            _xs.append('x{}'.format(i + 1))
        for i in range(nums):
            _subdic[_xs[i]] = x[i][0]
        for i in range(nums):
            _delta = diff(f_x, _xs[i])
            _delta_f.append(_delta.evalf(subs=_subdic))
        return np.array(_delta_f)

    def isInner(A1: np.ndarray, E: np.ndarray):
        '''
        :param A1,E:分离后的等式约束系数矩阵
        :return: A1,E是否都为零矩阵
        :testData:
            A1=np.zeros((5,5))
            E=np.zeros((3,5))
            print(ZouUtils.isInner(A1,E))
        '''
        return np.all(A1 == 0) and np.all(E == 0)

    def getMaxLambda(b2, A2, x, d):
        '''
        :param b2:分离后的右值矩阵
        :param A2:分离后的不等式约束矩阵
        :param x:当前的自变量
        :param d:可行下降方向
        :return:步长lambda
        '''
        _frac_up = b2 - np.matmul(A2, x)
        _frac_down = np.matmul(A2, d)
        _lambda = -1
        for i in range(len(_frac_down)):
            if _frac_down[i][0] >= 0:
                continue
            else:
                _lambda = max(_lambda, _frac_up[i][0] / _frac_down[i][0])
        if _lambda > -1:
            return _lambda
        return 0


class GradUtils:

    # 直接调用ZouUtils中的分割函数
    def saperateMatrix(A, x, b):
        return ZouUtils.saperateMatrix(A, x, b)

    def genProjection(A1:np.ndarray,E=np.empty((0,0))):
        '''
        :param A1: 有效约束矩阵
        :param E: 等式约束矩阵
        :return: 返回一个投影矩阵，是否有效不论
        '''
        _A_shape=A1.shape
        _E_shape=E.shape
        if(_A_shape==(0,0) and (_E_shape)==(0,0)):
            return np.identity(0)
        if _A_shape==(0,0):
            _M=E
        elif _E_shape==(0,0):
            _M=A1
        else:
            _M=np.append(A1,E,axis=0)
        _M_T=_M.T
        _MMT_inv=np.linalg.inv(np.matmul(_M,_M_T))
        _Q=np.matmul(_M_T,_MMT_inv)
        _Q=np.matmul(_Q,_M)
        _Q_shape=_Q.shape
        assert(_Q_shape[0] == _Q_shape[1])
        return np.identity(_Q_shape[0])-_Q

    def getGradient(f_x, x):
        '''
        :param f_x:目标函数的sympy表达式
        :param x: 自变量取值x.shape=n*1
        :return:梯度
        :testData:
            x=np.array([1,2]).reshape((2,1))
            x1, x2 = symbols('x1 x2')
            Y = x1 ** 2 - x2 ** 2
            print(ZouUtils.getGradient(Y,x))
        '''
        nums = max(x.shape)
        _xs = list()
        _subdic = dict()
        _delta_f = list()
        for i in range(nums):
            _xs.append('x{}'.format(i + 1))
        for i in range(nums):
            _subdic[_xs[i]] = x[i][0]
        for i in range(nums):
            _delta = diff(f_x, _xs[i])
            _delta_f.append(_delta.evalf(subs=_subdic))
        return np.array(_delta_f)

    def correctA1(A1,delta_f):
        '''
        :param A1: 待修正的有效约束矩阵
        :param M: A1和E的增广矩阵
        :param delta_f: 梯度
        :return: 修正之后的有效约束矩阵
        '''
        _isKTP=false
        _lambda = np.matmul(np.linalg.inv(np.matmul(A1, A1.T)), np.matmul(A1, delta_f.reshape(len(delta_f), 1)))
        _correctIndex = np.array(np.where(_lambda >= 0))[0]
        if np.all(_lambda>=0):
            _isKTP=true
        return {'A1':A1[_correctIndex],
                'KT':_isKTP}

    def getMaxLambda(A2,b2,x,d):
        '''
        :param b2:分离后的右值矩阵
        :param A2:分离后的不等式约束矩阵
        :param x:当前的自变量
        :param d:可行下降方向
        :return:步长lambda
        '''
        _frac_up = b2 - np.matmul(A2, x)
        _frac_down = np.matmul(A2, d)
        if np.all(_frac_down>=0):
            return inf
        _lambda = inf
        for i in range(len(_frac_down)):
            if _frac_down[i][0] >= 0:
                continue
            else:
                _lambda = min(_lambda, _frac_up[i][0] / _frac_down[i][0])
        if _lambda > -1:
            return _lambda
        return 0



if __name__ == '__main__':
    data = getData()

    f_x = data['f_x']
    A = data['A']
    E=data['E']
    print(E.shape)
    b = data['b']
    x = data['x']
    Saperate=GradUtils.saperateMatrix(A,x,b)
    print(Saperate)
    A1=Saperate['A1']
    A2=Saperate['A2']
    b2=Saperate['b2']
    Projection=GradUtils.genProjection(A1, E)
    print(Projection)
    delta_f_x=GradUtils.getGradient(f_x, x)
    print(delta_f_x)
    A1=GradUtils.correctA1(A1,delta_f_x)
    print(A1)
    A1=A1['A1']
    Projection=GradUtils.genProjection(A1)
    print(Projection)
    d=-np.matmul(Projection,delta_f_x.reshape(len(delta_f_x),1))
    print(d)
    maxLambda=GradUtils.getMaxLambda(A2,b2,x,d)
    print(maxLambda)
    l=symbols('l')
    xl=x+l*d
    xl=list2dict(xl)
    print(xl)
    f_l=f_x.evalf(subs=xl)
    print(f_l)
    print(getMinPoint(f_l,maxLambda,l))
    x=x+getMinPoint(f_l,maxLambda,l)*d
    print(x)


